import torch
import models_xh as models
from data_loader import prepare_image_cv2
import cv2
import numpy as np



# device = "cpu"
device = "cuda"
#必须跟训练时定义的网络一样
model = models.resnet80().to(torch.device(device))
# model.load_state_dict(torch.load(r'xhrs80_1024_0.pt', map_location='cpu'))
model.load_state_dict(torch.load(r'./good/xhrs80_j512_184.pt', map_location=device))
# model.load_state_dict(torch.load(r'./good/xhrs80_y512_886.pt', map_location='cuda'))
model.eval()
print(model)
std_size = 512
example = torch.rand(1, 3, std_size, std_size).to(torch.device(device))

traced_script_module = torch.jit.trace(model, example)

output = traced_script_module(example)
print(output[-1])
traced_script_module.save(r"seg_model_jin.pt")
# traced_script_module.save(r"seg_model_yin.pt")
# traced_script_module.save(r"xhrs80_1024_2760.pt")

